from flask import Flask, request, jsonify, render_template, session
from flask_cors import CORS
import requests
import sys
import markdown
import json
import os
from datetime import datetime
from dotenv import load_dotenv
import argparse

# Load environment variables
load_dotenv()

app = Flask(__name__)
CORS(app)
app.secret_key = 'your-secret-key-here'  # 用于 session 加密

# Ollama API configuration
OLLAMA_API_URL = os.getenv('OLLAMA_API_URL', 'http://localhost:11434')
MODEL = os.getenv('MODEL_NAME', 'llama2')
PORT = int(os.getenv('PORT', '5000'))

# 对话历史存储路径
CHAT_HISTORY_DIR = "chat_histories"
if not os.path.exists(CHAT_HISTORY_DIR):
    os.makedirs(CHAT_HISTORY_DIR)

# 配置存储路径
CONFIG_FILE = "config.json"

def load_config():
    """加载配置"""
    if os.path.exists(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, 'r', encoding='utf-8') as f:
                return json.load(f)
        except json.JSONDecodeError:
            return {}
    return {}

def save_config(config):
    """保存配置"""
    try:
        with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
        return True
    except Exception as e:
        print(f"Error saving config: {str(e)}")
        return False

def test_connection(api_url, token, model=None):
    """测试 Ollama API 连接"""
    try:
        # 检查 API URL 是否有效
        if not api_url.startswith(('http://', 'https://')):
            return False, 'Invalid API URL format. Must start with http:// or https://'
            
        # 检查 token 是否有效
        if not token:
            return False, 'Token is required'
            
        # 测试连接 - 使用 /api/status 端点来检查服务是否可用
        response = requests.get(
            f"{api_url}/api/status",
            headers={
                'Content-Type': 'application/json',
                'Authorization': token
            },
            timeout=10  # 设置超时时间为 10 秒
        )
        
        # 检查响应内容
        if response.status_code == 200:
            try:
                response_data = response.json()
                if response_data.get('status') == 'Ollama is running':
                    return True, 'Connection successful'
                else:
                    return False, 'Invalid response from server'
            except ValueError:
                return False, 'Invalid JSON response from server'
        elif response.status_code == 401:
            return False, 'Unauthorized - Invalid token. Please check if your token is correct.'
        elif response.status_code == 404:
            return False, 'API endpoint not found. Please check if the server is running and the API URL is correct.'
        elif response.status_code == 500:
            return False, 'Server error. The server might be experiencing issues. Please try again later.'
        elif response.status_code == 503:
            return False, 'Service unavailable. The server might be starting up or overloaded.'
        else:
            return False, f'Connection failed with status code {response.status_code}. Please check your configuration and try again.'
            
    except requests.exceptions.Timeout:
        return False, 'Connection timed out. Please check if the server is running and accessible.'
    except requests.exceptions.ConnectionError:
        return False, 'Could not connect to the server. Please check if the service is running and the API URL is correct.'
    except Exception as e:
        return False, f'Connection failed: {str(e)}. Please check your configuration and try again.'

def get_chat_history(chat_id):
    """获取指定对话的历史记录"""
    history_file = os.path.join(CHAT_HISTORY_DIR, f"{chat_id}.json")
    if os.path.exists(history_file):
        try:
            with open(history_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except json.JSONDecodeError:
            print(f"Error reading chat history for {chat_id}")
            return []
    return []

def save_chat_history(chat_id, messages):
    """保存对话历史记录"""
    try:
        history_file = os.path.join(CHAT_HISTORY_DIR, f"{chat_id}.json")
        with open(history_file, 'w', encoding='utf-8') as f:
            json.dump(messages, f, ensure_ascii=False, indent=2)
        return True
    except Exception as e:
        print(f"Error saving chat history: {str(e)}")
        return False

def get_all_chats():
    """获取所有对话列表"""
    chats = []
    try:
        for file in os.listdir(CHAT_HISTORY_DIR):
            if file.endswith('.json'):
                chat_id = file[:-5]  # 移除 .json 后缀
                history = get_chat_history(chat_id)
                if history:
                    # 获取最后一条消息的时间
                    last_updated = os.path.getmtime(os.path.join(CHAT_HISTORY_DIR, file))
                    # 使用第一条消息作为标题，如果没有则使用默认标题
                    title = "New Chat"
                    if history:
                        for msg in history:
                            if msg["role"] == "user":
                                title = msg["content"][:50] + '...' if len(msg["content"]) > 50 else msg["content"]
                                break
                    chats.append({
                        'id': chat_id,
                        'title': title,
                        'last_updated': last_updated
                    })
    except Exception as e:
        print(f"Error getting chat list: {str(e)}")
    return sorted(chats, key=lambda x: x['last_updated'], reverse=True)

def main():
    parser = argparse.ArgumentParser(description='Ollama Chat Web Interface')
    parser.add_argument('--host', default='0.0.0.0', help='Host to bind to')
    parser.add_argument('--port', type=int, default=5000, help='Port to bind to')
    parser.add_argument('--api-endpoint', default='http://localhost:11434', help='Ollama API endpoint')
    parser.add_argument('--token', default='', help='API token (optional)')
    parser.add_argument('--model', default='llama2', help='Model to use')
    args = parser.parse_args()

    # 初始化配置
    config = {
        'api_url': args.api_endpoint,
        'token': args.token,
        'model': args.model
    }
    
    # 保存配置
    save_config(config)
    
    # 启动服务器
    app.run(host=args.host, port=args.port, debug=True)

@app.route('/')
def index():
    config = load_config()
    chats = get_all_chats()
    current_chat_id = datetime.now().strftime('%Y%m%d%H%M%S')
    
    # 如果没有 token，显示配置提示
    show_config_prompt = not config.get('token')
    
    return render_template('index.html', 
                         api_endpoint=config['api_url'],
                         token=config['token'],
                         model=config['model'],
                         chats=chats,
                         current_chat_id=current_chat_id,
                         show_config_prompt=show_config_prompt)

@app.route('/api/chats', methods=['GET'])
def get_chats():
    return jsonify(get_all_chats())

@app.route('/api/chats/new', methods=['POST'])
def new_chat():
    try:
        chat_id = datetime.now().strftime('%Y%m%d%H%M%S')
        session['current_chat_id'] = chat_id
        
        # 创建空的对话历史文件
        history_file = os.path.join(CHAT_HISTORY_DIR, f"{chat_id}.json")
        with open(history_file, 'w', encoding='utf-8') as f:
            json.dump([], f)
            
        return jsonify({
            'chat_id': chat_id,
            'status': 'success',
            'message': 'New chat created successfully'
        })
    except Exception as e:
        print(f"Error creating new chat: {str(e)}")
        return jsonify({
            'status': 'error',
            'message': 'Failed to create new chat'
        }), 500

@app.route('/api/chats/<chat_id>', methods=['GET'])
def get_chat(chat_id):
    try:
        messages = get_chat_history(chat_id)
        return jsonify(messages)
    except Exception as e:
        print(f"Error getting chat history: {str(e)}")
        return jsonify({
            'status': 'error',
            'message': 'Failed to get chat history'
        }), 500

@app.route('/api/chats/<chat_id>/switch', methods=['POST'])
def switch_chat(chat_id):
    try:
        session['current_chat_id'] = chat_id
        return jsonify({
            'status': 'success',
            'message': 'Chat switched successfully'
        })
    except Exception as e:
        print(f"Error switching chat: {str(e)}")
        return jsonify({
            'status': 'error',
            'message': 'Failed to switch chat'
        }), 500

@app.route('/api/chats/<chat_id>/delete', methods=['DELETE'])
def delete_chat(chat_id):
    try:
        # 确保 chat_id 是有效的文件名
        if not chat_id or not chat_id.isalnum():
            return jsonify({
                'status': 'error',
                'message': 'Invalid chat ID'
            }), 400

        history_file = os.path.join(CHAT_HISTORY_DIR, f"{chat_id}.json")
        
        # 检查文件是否存在
        if not os.path.exists(history_file):
            print(f"Chat file not found: {history_file}")
            return jsonify({
                'status': 'error',
                'message': 'Chat not found'
            }), 404

        # 尝试删除文件
        try:
            os.remove(history_file)
            print(f"Successfully deleted chat: {chat_id}")
            return jsonify({
                'status': 'success',
                'message': 'Chat deleted successfully'
            })
        except OSError as e:
            print(f"Error deleting chat file: {str(e)}")
            return jsonify({
                'status': 'error',
                'message': 'Failed to delete chat file'
            }), 500

    except Exception as e:
        print(f"Error in delete_chat: {str(e)}")
        return jsonify({
            'status': 'error',
            'message': 'Internal server error'
        }), 500

@app.route('/api/chat', methods=['POST'])
def chat():
    config = load_config()
    if not config.get('token'):
        return jsonify({'error': 'Please configure your API token first'})
    
    data = request.get_json()
    message = data.get('message')
    
    if not message:
        return jsonify({'error': 'Message is required'})
    
    try:
        # 获取当前对话历史
        chat_id = session.get('current_chat_id')
        history = get_chat_history(chat_id)
        
        # 构建消息历史
        messages = [{"role": msg["role"], "content": msg["content"]} for msg in history]
        messages.append({"role": "user", "content": message})

        # 发送请求到 Ollama API
        response = requests.post(
            f"{config['api_url']}/api/chat",
            json={
                "model": config['model'],
                "messages": messages,
                "stream": False
            },
            headers={
                'Content-Type': 'application/json',
                'Authorization': config['token']  # 直接使用 token，不添加 'Bearer' 前缀
            }
        )

        if response.status_code == 200:
            response_data = response.json()
            content = response_data.get('message', {}).get('content', '')
            
            # 保存对话历史
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": content})
            if not save_chat_history(chat_id, history):
                print(f"Warning: Failed to save chat history for {chat_id}")
            
            return jsonify({'response': content})
        elif response.status_code == 401:
            return jsonify({'error': 'Unauthorized - Invalid token'})
        else:
            return jsonify({'error': 'Failed to get response from the model'})
    except Exception as e:
        return jsonify({'error': str(e)})

@app.route('/api/config', methods=['GET'])
def get_config():
    """获取当前配置"""
    config = load_config()
    return jsonify({
        'api_url': config.get('api_url', 'http://localhost:11434'),
        'token': config.get('token', ''),
        'model': config.get('model', 'llama2')
    })

@app.route('/api/config', methods=['POST'])
def update_config():
    """更新配置"""
    try:
        data = request.get_json()
        config = load_config()
        
        # 更新配置
        if 'api_url' in data:
            config['api_url'] = data['api_url']
        if 'token' in data:
            config['token'] = data['token']
        if 'model' in data:
            config['model'] = data['model']
        
        # 测试连接
        if 'api_url' in data or 'token' in data:
            success, message = test_connection(config.get('api_url'), config.get('token'), config.get('model'))
            if not success:
                return jsonify({
                    'status': 'error',
                    'message': message
                }), 400
        
        # 保存配置
        if save_config(config):
            return jsonify({
                'status': 'success',
                'message': 'Configuration updated successfully'
            })
        else:
            return jsonify({
                'status': 'error',
                'message': 'Failed to save configuration'
            }), 500
            
    except Exception as e:
        print(f"Error updating config: {str(e)}")
        return jsonify({
            'status': 'error',
            'message': 'Internal server error'
        }), 500

@app.route('/api/config/test', methods=['POST'])
def test_config():
    data = request.get_json()
    api_url = data.get('api_url')
    token = data.get('token')
    
    if not token:
        return jsonify({'status': 'error', 'message': 'Token is required'})
    
    success, message = test_connection(api_url, token)
    
    if success:
        # 如果测试成功，更新配置
        config = load_config()
        config['api_url'] = api_url
        config['token'] = token
        save_config(config)
        
    return jsonify({
        'status': 'success' if success else 'error',
        'message': message
    })

if __name__ == '__main__':
    main() 